﻿using Active.Core;        // Core API
using FlyIL.Expression;
using FlyIL.Misc;
using System.Collections.Generic;
using System.Linq;
using static FlyIL.Expression.PlaceHolder;
using static FlyIL.OperandMaker;
using static FlyIL.AssemblerRegisters;
using System;

namespace FlyIL.Trace;


internal class TraceCore
{

    static readonly AlterCase alter = new();



    internal bool CheckFactor { get; set; }


    Context Vcontext;


    /// <summary>
    /// 追踪的目标操作数 
    /// </summary>
    Operand Target;

    /// <summary>
    /// 要监视的中间操作数 
    /// </summary>
    List<Operand> Ops;

    /// <summary>
    /// 表达式
    /// </summary>
    ExpByTrace Main;


    /// <summary>
    /// 追踪的起始节点
    /// </summary>
    Inst From;

    /// <summary>
    /// 当前的节点
    /// </summary>
    Inst CNode { get => _cnode; set { _cnode = value; if (value != null) Last = value; } }
    internal Inst _cnode;


    public Inst Last { get; private set; }

    /// <summary>
    /// 当前的语句 
    /// </summary>
    Inst CInst { get => CNode; }



    /// <summary>
    /// !!! 记录追踪过程的， 栈的视角变化 , 用以在后期 恢复到 追踪语句的 视角 
    /// </summary>
    int AllSpEffect;


    #region 构造

    TraceCore(Context context, Inst node)
    {
        Vcontext = context;
        Ops = new();
        From = node;
    }
    internal TraceCore(Context context, Inst node, Operand target) : this(context, node)
    {

        Target = target;
        Ops = new() { target };
        Main = new(target);

    }


    /// <summary>
    /// 追踪一个表达式
    /// </summary>
    internal TraceCore(Context context, Inst node, Expr exp) : this(context, node)
    {
        var e = exp.Copy();
        Main = new ExpByTrace(e.Left, e.Opt, e.Right);

        /// 将表达式中的 RegCode 替换为 Operand
        Main.Recursion(o =>
        {
            if (o.Left is RegCode lcode)
                o.Left = new Operand(lcode);
            if (o.Right is RegCode rcode)
                o.Right = new Operand(rcode);
        });

        Ops = new();
        foreach (var u in Main.Collect<Operand>())
            Ops.CheckAdd(u);

    }

    #endregion


    internal ExpByTrace Trace(Inst stop)
    {
        var ori = Main;

        for (CNode = From;
            Ops.Count > 0 && CNode != null && CNode != stop;
            CNode = Vcontext.GetPreviousIfExist(CNode))
        {

            var sta = CheckUnreachable() || CheckEndPrefix() ||
                        (
                            CheckBitFlags()
                            + CheckExpand()
                            + CheckFolder()
                            + FixStackOffset()
                            + (CheckCutOff() || CheckAlter())
                            + CheckExpCutOff()
                        );

            ;

        }


        if (Main is null)
            return null;


        /*****************************************************
         * 
         * 要确保表达式的原子性 
         * 如果追踪位置 还在原来的节点上，那么直接返回null 即可
         * 
         ***********************************/

        if (Main.Depth == 1)
            if (Main.Ext == From)
                return null;


        /********************
         * 检测有没有 直接取址
         * *********************/
        Main.Recursion(o =>
        {
            if (o.Left is Operand oo && oo.IsConstHeap)
            {
                o.Left = oo.PtrSize switch
                {
                    1 => Vcontext.Reader.Read8(oo.u64),
                    2 => Vcontext.Reader.Read16(oo.u64),
                    4 => Vcontext.Reader.Read32(oo.u64),
                    _ => (object)Vcontext.Reader.Read64(oo.u64),
                };
            }
        });


        /// 对追踪到的 栈变量 修正回到 追踪起始位置的 偏移值  
        Main?.Recursion((o) =>
        {
            if (o.Left is Operand lo)
                if (lo.IsStackValue)
                {
                    lo.u32 = (uint)((int)(lo.u32) - AllSpEffect);
                    o.Left = lo;
                }

            if (o.Right is Operand ro)
                if (ro.IsStackValue)
                {
                    ro.u32 = (uint)((int)(ro.u32) - AllSpEffect);
                    o.Right = ro;
                }

        });




        /**********
         * 追踪到 分段 赋值 (两个操作数 ) ， 需要手动合并  
         * ******/

        var a = new Expr(Any1, Opt.mov, _Op1).mov(_Op2);
        var b = new Expr(new Expr(Any1), Opt.mov, _Op1).mov(_Op2);
        var c = new Expr(_Op1).mov(_Op2);

        Main.Recursion((o) =>
        {
            /// 形如 ( (0 <= a ) <= b )
            if (a.IsMatch(o) || b.IsMatch(o))
            {
                var lo = (o.Left as Expr).FinalRight<Operand>();
                var ro = o.FinalRight<Operand>();

                var newo = lo + (ro);
                if (newo != null)
                {
                    (o.Left as Expr).Right = newo;
                    o.Opt = Opt.nop;
                    o.Right = null;
                }
            }
            /// 形如 ( (a) <= b)
            else if (c.IsMatch(o))
            {
                var a = (o).FinalLeft<Operand>();
                var b = o.FinalRight<Operand>();

                var newo = a + (b);
                if (newo != null)
                {
                    (o.Left as Expr).Left = newo;
                    o.Opt = Opt.nop;
                    o.Right = null;
                }
            }
        });



        Main.Recursion((o) =>
        {
            /*********
             * 对无效的 load 节点 进行回退 ， 比如  load? [  ptr[esi] + ptr [eax] ] 
             * *************/
            if (o.Opt.IsLoad())
            {

                var ptr = o.Find<Operand?>(o => o != null && (bool)o?.IsPtr);
                if (ptr != null)
                {
                    (o).Replace(ptr, o.Ext.O1);
                }
            }
            /**********
             * 检测 表达式中的  load? 节点 ，尝试将它转为操作数 
             * ******/
            if (o.Opt.IsLoad())
            {
                var op = Operand.FromLoadExp(ref o);
                if (op is null)
                    return;

                o.Left = op;
                o.Opt = Opt.nop;

            }

        });

        return Main;
    }


    #region 追踪逻辑

    /// <summary>
    /// 以下 寄存器 无意义 ， 比如 eflag  或者 rsp , esp (??? 不是取址的情况下 )
    /// </summary>
    /// <returns> 返回done 使循环退出 </returns>
    status CheckUnreachable()
    {

        if (Ops.Any(o => o.IsReg(RegCode.qflag, RegCode.eflag, RegCode.flag)))
        {
            CNode = null;
            return status._done;
        }

        // rsp 或者 
        if (Ops.Any(o => o.IsRegister && o.Base.LE(RegCode.rsp)))
        {
            CNode = null;

            // 甚至 sp 以下的， 均无意义
            if (Ops.Any(o => o.IsRegister && o.Base.LE(RegCode.sp)))
                Main = null;

            return status._done;
        }

        var node = Main.Find(o => o.Left is Operand lo && lo.IsReg(RegCode.esp, RegCode.rsp) || o.Right is Operand ro && ro.IsReg(RegCode.esp, RegCode.rsp));
        if (null != node)
        {
            var load = node.UnperFind(o => o.Opt.IsLoad());
            if (null == load)
            {

                CNode = null;
                return status._done;
            }
        }



        return status._fail;
    }

    /// <summary>
    /// 
    /// </summary>
    /// <returns>返回done 使循环退出</returns>
    /// <exception cref="System.Exception"></exception>
    status CheckEndPrefix()
    {

        if (CInst.Prefix != Prefix.End)
            return status._fail;

        /// ecx 是rep 指令的计数器 
        var ecx = _regs_[R3D];
        if (Ops.HaveGe(ecx))
            return status._fail;

        var start = CNode;
        while (start != null && !start.HavePrefix)
        {
            if (start == null)
                throw new Exception("未找到匹配的 rep 头");
            start = Vcontext.GetPreviousIfExist(start);
        }
        var list = Vcontext.CollectOperands(start, CNode);


        bool found = false;
        foreach (var li in list)
        {
            var _ = li;
            /// 如果存在 Base , 也需要退出 
            if (Ops.HaveGe(_) || Ops.HaveBase(_))
            {
                found = true;
                break;
            }
        }

        if (found) return status._fail;

        CNode = Vcontext.GetPreviousIfExist(start);

        return status._done;
    }

    status pushCase()
    {
        var inst = CNode;

        /// 这里需要循环，一个push 可能对应多个目标 
        for (int i = Ops.Count - 1; i >= 0; i--)
        {
            /// 追踪的是一个 栈内变量 
            if (!Ops[i].IsStackValue)
                continue;

            // 偏移符合要求
            if ((int)Ops[i].u32 < 0 || Ops[i].u32 > inst.O1.ByteSize)
                continue;

            if (Ops[i].u32 + Ops[i].ByteSize <= inst.O1.ByteSize)
            {

                Operand old = Ops[i];
                Ops.RemoveAt(i);

                Operand op = inst.O1;


                /// 来自于 堆内存 
                if (inst.O1.IsHeapValue)
                {
                    /// 如果是堆地址 + 偏移 或者 有Index  的形式 ， 不再追踪 
                    //if (inst.Op0.u32 != 0 || inst.Op0.Index != None)
                    //{
                    //    return;
                    //}

                    var op_ = new Operand(inst.O1.Base);
                    if (!Ops.HaveGe(op))
                    {
                        Ops.CheckAdd(inst.O1);

                        Main.Replace(old, new Expr(inst.O1, CNode));
                    }
                }
                else
                {
                    /// 压栈与弹栈是对称的
                    if (old.BitSize == inst.O1.BitSize)
                    {
                        if (inst.Opt == Opt.call)
                        {
                            var o_ = Operand.Make(inst.Vip + (uint)inst.O3, Vcontext.Is64 ? 8 : 4);

                            Main.Replace(old, new Expr { Left = o_, Ext = CNode });
                        }
                        else
                        {
                            /// 因为这里是一个 push ， 而后面要修改栈便宜，所以加一个提前量 
                            if (op.IsStackValue)
                                op.u32 += (uint)op.BitSize / 8;

                            Ops.CheckAdd(op);

                            Main.Replace(old, new Expr { Left = op, Ext = CNode });
                        }
                    }
                    /// 不对称 
                    else
                    {
                        /// 如果是寄存器
                        if (inst.O1.IsRegister)
                        {
                            var op_ = new Operand(inst.O1.Base.ChangTo(old.ByteSize));
                            Ops.CheckAdd(op_);
                            Main.Replace(old, new Expr { Left = op_, Ext = CNode });
                        }
                        /// 此处不再需要追踪了 
                        else if (inst.O1.IsImm)
                        {
                            int size = old.ByteSize;
                            ulong a = size == 8 ? inst.O1.u32 : size == 4 ? inst.O1.u16 : inst.O1.u8;
                            Main.Replace(old, new Expr { Left = Operand.Make(a, size), Ext = CNode });
                        }
                    }
                }

            }
            else
            {
                return status._fail;
            }

        }
        return status._done;
    }

    status FixStackOffset()
    {
        int? spIndex = CInst.SpVariety;

        if (spIndex == null || spIndex == 0)
            return status._fail;
        //throw new Exception("无法计算栈深度");

        AllSpEffect += (int)spIndex;


        /// 入栈的追踪, 需要 特殊处理 
        if (CInst.Opt == Opt.push || CInst.Opt == Opt.call)
            pushCase();

        Main.Recursion((exp) =>
        {
            if (exp.Left is Operand le && le.IsStackValue)
            {
                le.u32 += (uint)spIndex;
                exp.Left = le;
            }
            if (exp.Right is Operand re && re.IsStackValue)
            {
                re.u32 += (uint)spIndex;
                exp.Right = re;
            }
        });

        /// 如果栈有变化 ，修正 要监视的 栈操作数 
        Ops.FixSp((int)spIndex);

        if (Target.IsStackValue)
            Target.u32 += (uint)spIndex;

        return status._fail;
    }

    /// <summary>
    /// 处理标志位
    /// </summary>
    status CheckBitFlags()
    {

        /// 先检测 标志位 
        if (Ops.HaveBitFlags(CInst.Opt))
        {
            //if (!Vcontext.OptimizerOption.TraceBitFlags)
            //    return status._done;

            /// 处理表达式
            var flags = Ops.GetBitFlags(CInst.Opt);
            foreach (var u in flags)
            {
                Main.Recursion((exp) =>
                {
                    if (exp.Left is Operand lo && lo == u)
                    {
                        exp.Left = new Expr(new Expr(CInst.O1, CInst.Opt, CInst.OpsCount > 1 ? CInst.O2 : null), u.Base.GetIf(), null);
                    }
                });
            }

            /// 清理追踪目标 
            Ops.RemoveBitFlags(CInst.Opt);
            if (!Ops.HaveGe(CInst.O1))
                Ops.CheckAdd(CInst.O1);

            if (CInst.OpsCount > 1)
                if (!Ops.HaveGe(CInst.O2))
                    Ops.CheckAdd(CInst.O2);



        }

        return status._done;
    }

    /// <summary>
    /// 对语句进行 截止判断 
    /// 如 xor eax,eax , 则不再需要追踪 ax的值 
    /// </summary>
    status CheckCutOff()
    {
        if (Ops.HaveLe(CInst.O1) && CInst.MatchIn(_Op1.mov(0), _Op1 ^ _Op1, _Op1 ^ 0, _Op1 - _Op1))
        {
            Ops.RemoveLe(CInst.O1);

            Main.Recursion((exp) =>
            {
                if (CInst.O1.Equals(exp.Left))
                {
                    if (exp.Opt == Opt.mov)
                    {
                        exp.Opt = Opt.nop;
                        exp.Left = exp.Right;
                        exp.Right = null;
                    }
                    else
                        exp.Left = (ulong)0;

                }
                if (CInst.O1.Equals(exp.Right))
                    exp.Right = (ulong)0;

            });

            /// 如果没有 需要监视的 操作数 , 则视为起点 
            if (Ops.Count == 0)
                return status._fail;

            //    break;
            //continue;

            return status._done;
        }


        return status._fail;

    }

    /// <summary>
    /// 对表达式 进行截止 判断 
    /// </summary>
    status CheckExpCutOff()
    {

        var o = Main.Find(o => (Any1.cmp(Any1)).IsMatch(o));

        if (o != null)
        {
            Main.Replace(o, new Expr(1, Opt.cmp, 1));

            Ops.Clear();
            Ops.AddRange(Main.Collect<Operand>());
        }


        return status._fail;
    }

    /// <summary>
    /// 检测如果有 取址操作数 ， 且与现在的 操作数有交集， 则需要展开 
    /// </summary>
    status CheckExpand()
    {
        if (Ops.HavePtrAnd(CInst.O1))
        {
            Ops.Expand(ref CInst.O1);

            Main.Recursion((exp) =>
            {
                if (exp.Left is Operand lo && lo.IsPtr)
                {
                    if (lo.Base.And(CInst.O1.Base))
                        exp.Left = lo.GetExp();
                    if (lo.Index.And(CInst.O1.Base))
                        exp.Left = lo.GetExp();

                    (exp.Left as Expr).Ext = CNode;
                }
                if (exp.Right is Operand ro && ro.IsPtr)
                {
                    if (ro.Base.And(CInst.O1.Base))
                        exp.Left = ro.GetExp();
                    if (ro.Index.And(CInst.O1.Base))
                        exp.Left = ro.GetExp();
                    (exp.Left as Expr).Ext = CNode;
                }

            });
            return status._done;
        }


        return status._fail;
    }

    /// <summary>
    /// 检查 表达式中 是否有 栈变量, 并收缩回 操作数
    /// </summary>
    status CheckFolder()
    {

        /// 如果 中间数中 有 esp , 那么需要 将表达式中的栈变量, 转回操作数, 以移除中间数中的 esp
        if (!Ops.HaveEq(new Operand(Vcontext.CodeSp)))
            return status._fail;


        var loadStack = Main.Find(exp =>  exp.Opt.IsLoad() && exp.FinalLeft<Operand>().IsReg(RegCode.rsp, RegCode.esp) );

        /// &4 (esp +4) 
        if (loadStack != null)
        {
            var op = Operand.FromLoadExp(ref loadStack);
            if (op is null)
                return status._fail;


            Main.Replace(loadStack, op);

            Ops.Remove(new(Vcontext.CodeSp));
            Ops.CheckAdd((Operand)op);
        }

        return status._fail;
    }

    status CheckAlter()
    {

        if (!CInst.IsAlter)
            return status._done;

        var pass = alter.GetSubMatch(CInst, Ops);

        if (pass == null)
        {
            if (CNode != From)
            {
                /// 消费了目标值 ， 且不是作用于 中间数
                var a = !Ops.HaveLe(CInst.O1);
                Main.HasTargetBeenUsed |= (a && CInst.Used(Target));
                Main.HasTargetBeenChanged |= CInst.WillAlter(Target);
            }
        }
        else
        {
            pass.Do(CNode, Main, Ops);
        }

        return status._done;
    }

    #endregion

}
